package com.tutou.langchain4j.config;

import com.tutou.langchain4j.service.Assistant;
import dev.langchain4j.data.segment.TextSegment;
import dev.langchain4j.memory.chat.MessageWindowChatMemory;
import dev.langchain4j.model.chat.ChatLanguageModel;
import dev.langchain4j.rag.content.retriever.EmbeddingStoreContentRetriever;
import dev.langchain4j.store.embedding.EmbeddingStore;
import dev.langchain4j.store.embedding.inmemory.InMemoryEmbeddingStore;
import dev.langchain4j.web.search.WebSearchTool;
import dev.langchain4j.web.search.searchapi.SearchApiWebSearchEngine;
import lombok.RequiredArgsConstructor;
import dev.langchain4j.service.AiServices;
import org.springframework.context.annotation.Bean;
import org.springframework.context.annotation.Configuration;

@Configuration
@RequiredArgsConstructor
public class AssistantInit {
    final ChatLanguageModel chatLanguageModel;

    /**
     * 初始化向量数据库工具
     * @return
     */
    @Bean
    public EmbeddingStore<TextSegment> initEmbeddingStore() {
        // 企业产品需要将内存向量数据库替换成更成熟的产品 PgVector，es,等等
        return new InMemoryEmbeddingStore<>();
    }

    /**
     * 初始化联网工具
     * @return
     */
    @Bean
    public WebSearchTool initWebSearch() {
        SearchApiWebSearchEngine searchEngine = SearchApiWebSearchEngine.builder()
                .apiKey("rSSTTfDWLecgbSjPhjSwCsmn")
                .engine("baidu")
                .build();
        return WebSearchTool.from(searchEngine);
    }

    @Bean
    public Assistant init(EmbeddingStore<TextSegment> embeddingStore, WebSearchTool webSearchTool) {
        return AiServices.builder(Assistant.class)
                // 会话记忆
                .chatMemoryProvider(memoryid -> MessageWindowChatMemory.withMaxMessages(10))
                // 模型配置
                .chatLanguageModel(chatLanguageModel)
                // 外挂内存向量属于库
                .contentRetriever(EmbeddingStoreContentRetriever.from(embeddingStore))
                // function call 设置
                .tools(new HighLeverCalculator(), webSearchTool)
                .build();
    }
}
